# -*- coding: utf-8 -*-
# @Time    : 2025/10/5
# @Author  : 楚江涛
# @FileName: executor.py
# @Software: PyCharm

"""
异步工作流执行器（修正版）
---------------------------------
修复点：
1. 将 asyncio.Semaphore / Event 等对象延迟到 run() 内初始化，
   以避免 “There is no current event loop in thread …” 错误；
2. 允许 run() 前调用 pause()/cancel() 等方法；
3. 保留原全部调度逻辑与事件上报机制。
"""

from __future__ import annotations
import asyncio
import json
import traceback
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Set

from domain.enums import EventType, ExecStatus, NodeType
from execution.interfaces import EventPublisher, CodeRunner


@dataclass(frozen=True)
class NodeCtx:
    id: str
    type: str
    data: Dict[str, Any]


class WorkflowExecutor:
    def __init__(
        self,
        exec_id: str,
        wf: Dict[str, Any],
        events: EventPublisher,
        runner: CodeRunner,
        *,
        max_parallelism: int = 8,
        node_timeout: Optional[float] = None,
    ) -> None:
        self.exec_id = exec_id
        self.wf = wf
        self.events = events
        self.runner = runner
        self.max_parallelism = max_parallelism
        self.node_timeout = node_timeout

        # ---------- 构建拓扑 ----------
        self.nodes: Dict[str, NodeCtx] = {
            n["id"]: NodeCtx(n["id"], n["type"], n.get("data") or {})
            for n in wf.get("nodes", [])
        }
        self.succ: Dict[str, List[str]] = {}
        self.pred: Dict[str, List[str]] = {}
        self.indeg: Dict[str, int] = {nid: 0 for nid in self.nodes}
        for e in wf.get("edges", []):
            s, t = e["sourceNodeID"], e["targetNodeID"]
            self.succ.setdefault(s, []).append(t)
            self.pred.setdefault(t, []).append(s)
            self.indeg[t] += 1

        # ---------- 运行期缓存 ----------
        self.node_inputs: Dict[str, Dict[str, Any]] = {}
        self.node_outputs: Dict[str, Dict[str, Any]] = {}

        # ---------- 调度/控制（延迟初始化） ----------
        self.sem: Optional[asyncio.Semaphore] = None
        self.cancel_event: Optional[asyncio.Event] = None
        self.pause_event: Optional[asyncio.Event] = None
        self._cancel_requested = False
        self._pause_requested = False

        self.started: Set[str] = set()
        self.running: Set[str] = set()
        self.finished: Set[str] = set()
        self.node_tasks: Dict[str, asyncio.Task] = {}

        self.failed = False

    # ============================= 外部控制 =============================
    def request_cancel(self) -> None:
        self._cancel_requested = True
        if self.cancel_event:
            self.cancel_event.set()
        self._emit(EventType.EXEC_CANCELLING, {"reason": "用户主动取消"})
        for nid, t in list(self.node_tasks.items()):
            if not t.done():
                t.cancel()

    def request_pause(self) -> None:
        self._pause_requested = True
        if self.pause_event and not self.pause_event.is_set():
            self.pause_event.set()
        self._emit(EventType.EXEC_PAUSED, {})

    def request_resume(self) -> None:
        self._pause_requested = False
        if self.pause_event and self.pause_event.is_set():
            self.pause_event.clear()
        self._emit(EventType.EXEC_RESUMED, {})

    def cancel_node(self, node_id: str) -> None:
        t = self.node_tasks.get(node_id)
        if t and not t.done():
            t.cancel()
            self._emit(EventType.NODE_CANCELLED, {"node_id": node_id})

    # ============================= 主执行 =============================
    async def run(self, initial_context: Optional[Dict[str, Any]] = None) -> None:
        # 在协程上下文内初始化 asyncio 对象
        if self.sem is None:
            self.sem = asyncio.Semaphore(self.max_parallelism)
        if self.cancel_event is None:
            self.cancel_event = asyncio.Event()
        if self.pause_event is None:
            self.pause_event = asyncio.Event()
        if self._cancel_requested:
            self.cancel_event.set()
        if self._pause_requested:
            self.pause_event.set()

        print(f"[Executor] 开始执行工作流 {self.exec_id}")
        initial_context = initial_context or {}

        # --------- Start 节点预处理 ---------
        start_nodes: List[str] = []
        for nid, n in self.nodes.items():
            if n.type == NodeType.START.value:
                self.node_outputs[nid] = n.data.get("outputs", {}) or {}
                start_nodes.append(nid)
        for s in start_nodes:
            for nxt in self.succ.get(s, []):
                self.indeg[nxt] -= 1

        ready: List[str] = [
            nid for nid, d in self.indeg.items()
            if d == 0 and self.nodes[nid].type != NodeType.START.value
        ]
        if not ready and all(self.nodes[nid].type == NodeType.START.value for nid in self.nodes):
            self._emit(EventType.EXEC_STARTED, {"ready": []})
            self._emit(EventType.EXEC_FINISHED, {"status": ExecStatus.SUCCEEDED.value})
            return
        if not ready and any(self.nodes[nid].type != NodeType.START.value for nid in self.nodes):
            self._emit(EventType.EXEC_STARTED, {"ready": []})
            self._emit(EventType.EXEC_FINISHED,
                       {"status": ExecStatus.FAILED.value, "error": "无可就绪节点：可能存在环或依赖错误"})
            return

        self._emit(EventType.EXEC_STARTED, {"ready": list(ready)})
        accounted_finished: Set[str] = set()

        try:
            while (ready or self.running) and not self.failed:
                while self.pause_event.is_set() and not self.cancel_event.is_set():
                    await asyncio.sleep(0.05)

                if self.cancel_event.is_set():
                    self.failed = True
                    break

                # 启动 ready 节点
                for nid in list(ready):
                    if self.cancel_event.is_set():
                        break
                    if nid in self.started or nid in self.finished:
                        continue
                    self.node_inputs[nid] = self._build_inputs_for(nid, initial_context)
                    t = asyncio.create_task(self._run_node(nid))
                    self.node_tasks[nid] = t
                    self.started.add(nid)
                    self.running.add(nid)
                    ready.remove(nid)

                if not self.node_tasks and not ready:
                    break

                if self.node_tasks:
                    done, _ = await asyncio.wait(
                        list(self.node_tasks.values()),
                        return_when=asyncio.FIRST_COMPLETED,
                    )
                    for t in done:
                        if t.exception():
                            self.failed = True
                    if self.failed:
                        for t in list(self.node_tasks.values()):
                            if not t.done():
                                t.cancel()
                        break

                newly = self.finished - accounted_finished
                if newly:
                    for nid in newly:
                        for nxt in self.succ.get(nid, []):
                            self.indeg[nxt] -= 1
                            if (
                                self.indeg[nxt] == 0
                                and nxt not in self.finished
                                and nxt not in self.running
                                and nxt not in self.started
                                and self.nodes[nxt].type != NodeType.START.value
                            ):
                                ready.append(nxt)
                    accounted_finished |= newly

            status = (
                ExecStatus.CANCELLED.value
                if self.cancel_event.is_set()
                else (ExecStatus.FAILED.value if self.failed else ExecStatus.SUCCEEDED.value)
            )
            self._emit(EventType.EXEC_FINISHED, {"status": status})

        except Exception as e:
            self.failed = True
            self._emit(
                EventType.EXEC_FINISHED,
                {"status": ExecStatus.FAILED.value, "error": str(e), "traceback": traceback.format_exc()[-1500:]},
            )

    # ============================= 节点执行 =============================
    async def _run_node(self, node_id: str) -> None:
        node = self.nodes[node_id]
        self._emit(EventType.NODE_STARTED, {"node_id": node_id, "type": node.type})
        try:
            if self.cancel_event.is_set():
                raise asyncio.CancelledError()

            async with self.sem:
                if node.type == NodeType.CODE.value:
                    result = await self._run_code_node(node_id, node)
                elif node.type == NodeType.END.value:
                    ins = self.node_inputs.get(node_id, {})
                    result = {"result": ins.get("result") or json.dumps(ins, ensure_ascii=False)[:2000]}
                elif node.type == NodeType.START.value:
                    result = node.data.get("outputs", {}) or {}
                else:
                    result = {}

            self.node_outputs[node_id] = result
            self.finished.add(node_id)
            self.running.discard(node_id)
            self.node_tasks.pop(node_id, None)
            self._emit(EventType.NODE_FINISHED, {"node_id": node_id, "result": result})

        except asyncio.CancelledError:
            self.running.discard(node_id)
            self.finished.add(node_id)
            self.node_tasks.pop(node_id, None)
            self._emit(EventType.NODE_CANCELLED, {"node_id": node_id})
            raise
        except Exception as e:
            self.running.discard(node_id)
            self.finished.add(node_id)
            self.node_tasks.pop(node_id, None)
            self.failed = True
            self._emit(
                EventType.NODE_FAILED,
                {"node_id": node_id, "error": str(e), "traceback": traceback.format_exc()[-1500:]},
            )

    async def _run_code_node(self, node_id: str, node: NodeCtx) -> Dict[str, Any]:
        codes = node.data.get("codes") or {}
        code = codes.get("value") or codes.get("default") or ""
        params = self.node_inputs.get(node_id, {})
        timeout = self.node_timeout
        return await self.runner.run(
            node_id=node_id,
            code=code,
            params=params,
            timeout=timeout,
            env=None,
            metadata=None,
        )

    # ============================= 输入构建与事件 =============================
    def _build_inputs_for(self, node_id: str, global_ctx: Dict[str, Any]) -> Dict[str, Any]:
        ins: Dict[str, Any] = dict(global_ctx or {})
        for p in self.pred.get(node_id, []):
            out = self.node_outputs.get(p)
            if isinstance(out, dict):
                ins.update(out)
            elif out is not None:
                ins[p] = out
        node = self.nodes[node_id]
        if node.data.get("inputsValues"):
            ins.update(node.data["inputsValues"])
        return ins

    def _emit(self, event: EventType, payload: Dict[str, Any]) -> None:
        payload = dict(payload or {})
        payload["exec_id"] = self.exec_id
        try:
            self.events.publish(self.exec_id, event, payload)
        except Exception:
            print("[Executor] 事件上报失败（忽略）")
            pass
